{-# LANGUAGE TypeOperators, TypeFamilies, TemplateHaskell,
  UndecidableInstances, FlexibleInstances, MultiParamTypeClasses,
  ScopedTypeVariables, FlexibleContexts #-}

{-# OPTIONS_GHC -fcontext-stack=200 #-}

module Examples where

import TypeBasics
import Yoko
import Vector



data Wtree a = Wleaf  a
               | Wnode  (Wtree a) (Wtree a)
               | Weight (Wtree a) Double deriving Show

-- all of these declarations are to be generated
data Wleaf_  a = Wleaf_  a
data Wnode_  a = Wnode_  (Wtree a) (Wtree a)
data Weight_ a = Weight_ (Wtree a) Double
concat `fmap` mapM derive [''Wleaf_, ''Wnode_, ''Weight_]
type instance Range (Wleaf_  a) = Wtree a
type instance Range (Wnode_  a) = Wtree a
type instance Range (Weight_ a) = Wtree a
instance (True ~ Equal a a) => DC (Wleaf_  a) where
  rejoin (Wleaf_  a)     = Wleaf  a
instance (True ~ Equal a a) => DC (Wnode_  a) where
  rejoin (Wnode_  tl tr) = Wnode  tl tr
instance (True ~ Equal a a) => DC (Weight_ a) where
  rejoin (Weight_ t w)   = Weight t  w
type instance DCs (Wtree a) =
  N (Wleaf_ a) :+ N (Wnode_ a) :+ N (Weight_ a)
instance (True ~ Equal a a) => DT (Wtree a) where
  disband (Wleaf  a)     = disbanded $ Wleaf_  a
  disband (Wnode  tl tr) = disbanded $ Wnode_  tl tr
  disband (Weight t  w)  = disbanded $ Weight_ t  w
type instance Rep (Wleaf_  a) = D a
type instance Rep (Wnode_  a) = R (Wtree a) :* R (Wtree a)
type instance Rep (Weight_ a) = R (Wtree a) :* D Double
instance Generic (Wleaf_ a) where
  rep ~(Wleaf_ a) = D a
  obj ~(D a)      = Wleaf_ a
instance Generic (Wnode_ a) where
  rep ~(Wnode_ tl tr) = R tl :* R tr
  obj ~(R tl :* R tr) = Wnode_ tl tr
instance Generic (Weight_ a) where
  rep ~(Weight_ t w) = R t :* D w
  obj ~(R t :* D w)  = Weight_ t w





class GSize a where gsize :: a -> Int

instance GSize a => GSize (N a) where gsize = foldN gsize
instance (GSize a, GSize b) => GSize (a :+ b) where
  gsize = foldPlus gsize gsize

instance GSize U where gsize _ = 0
instance (GSize a, GSize b) => GSize (a :* b) where
  gsize (x :* y) = gsize x + gsize y

instance GSize a => GSize (D a) where gsize (D a) = gsize a
instance GSize a => GSize (R a) where gsize (R a) = gsize a



gsize_Wtree_0 = go where
  go t = case t of
    Wleaf  a     -> 1
    Wnode  tl tr -> go tl + go tr
    Weight t _   -> go t

gsize_Wtree_1 t = case project $ disband t of
  Left  (Wleaf_ a) -> 1
  Right x          -> case project x of
    Left  (Weight_ t w) -> gsize t
    Right x             -> gsize $ reps x

gsize_Wtree_2 t = case partition $ disband t of
  Left  x -> ($ x) $
    (\(Wleaf_ a) -> 1)   .|.   (\(Weight_ t w) -> gsize t)
  Right x -> gsize $ reps x

instance (True ~ Equal a a) => GSize (Wtree a) where
  gsize = gsize_Wtree_2 -- equivalently gsize_Wtree_1
-- NB equivalently
-- instance GSize (Wtree a) where gsize = gsize_Wtree_0




data Tree a = Leaf a | Node (Tree a) (Tree a) [Tree a] deriving Show
class ToTree el a where toTree_ :: a -> Tree el

instance ToTree el a => ToTree el (N a) where toTree_ = foldN toTree_
instance (ToTree el a, ToTree el b) => ToTree el (a :+ b) where
  toTree_ = foldPlus toTree_ toTree_

instance (ToTreeVec el (a :* b), Vec2Node (CountRs (a :* b))
         ) => ToTree el (a :* b) where toTree_ = vec2Node . toTrees_



class Vec2Node n where vec2Node :: Vector (Tree el) n -> Tree el
-- NB there is no legitimate instance for Z
instance Vec2Node (S Z)     where vec2Node = vHead
instance Vec2Node (S (S n)) where
  vec2Node v = Node (vHead v) (vHead v') $ vToList $ vTail v'
    where v' = vTail v



class ToTreeVec el prod where
  toTrees_ :: prod -> Vector (Tree el) (CountRs prod)

instance ToTreeVec el (D a) where toTrees_ _ = VNil
instance ToTreeVec el U     where toTrees_ _ = VNil

instance ToTree el a => ToTreeVec el (R a) where
  toTrees_ (R a) = VCons (toTree_ a) VNil

instance (ToTreeVec el a, ToTreeVec el b) => ToTreeVec el (a :* b) where
  toTrees_ (a :* b) = toTrees_ a `vAppend` toTrees_ b





type family El a
toTree :: ToTree (El a) a => a -> Tree (El a)
toTree = toTree_





toTree_Wtree_0 :: Wtree a -> Tree a
toTree_Wtree_0 = go where
  go t = case t of
    Wleaf  a     -> Leaf a
    Wnode  tl tr -> Node (go tl) (go tr) []
    Weight t  w  -> go t

{- ill-typed: "Could not deduce (Vec2Node Z)", because Wleaf has no Rs
toTree_Wtree_BAD ::
  (True ~ Equal a a, ToTree a (Wtree a)) => Wtree a -> Tree a
toTree_Wtree_BAD t = case t of
  Wleaf a -> Leaf a
  x       -> toTree_ . reps . disband $ x -}

toTree_Wtree_1 :: forall a.
  (True ~ Equal a a, ToTree a (Wtree a)) => Wtree a -> Tree a
toTree_Wtree_1 t = case project $ disband t of
  Left  (Wleaf_ a) -> Leaf a
  Right x          -> toTree_ $ reps x

instance (True ~ Equal a a) => ToTree a (Wtree a) where
 toTree_ = toTree_Wtree_1
-- NB equivalent
-- instance ToTree (Wtree a) a where toTree_ = toTree_Wtree_0



type instance El (Wtree a) = a





wt0 = Weight (Wleaf 'o') 1   `Wnode`   Weight (Wleaf 'k') 2

ex0 = gsize wt0
ex1 = toTree wt0





data ElT
type instance Apply ElT a = El a
